use image::GrayImage;
use tch::{Kind, Tensor};

fn tensor_to_byte_vec(tensor: &Tensor) -> Vec<u8> {
    let size = tensor.numel() * tensor.kind().elt_size_in_bytes();
    let mut data = vec![0u8; size];
    tensor.copy_data_u8(&mut data, tensor.numel());
    data
}

pub fn save_gray_28x28(tensor: &Tensor) {
    // 确保张量类型为f32
    assert_eq!(
        tensor.kind(),
        Kind::Float,
        "Tensor must be of type f32"
    );
    let element_size = tensor.kind().elt_size_in_bytes();
    assert_eq!(element_size, 4, "Element size must be 4 bytes for f32");

    // 确保张量尺寸为28x28
    let numel = tensor.numel();
    assert_eq!(numel, 28 * 28, "Tensor must have 28x28 elements");

    let byte_vec = tensor_to_byte_vec(tensor);
    assert_eq!(
        byte_vec.len(),
        numel * element_size,
        "Byte vector length mismatch"
    );

    // 检查内存对齐
    assert_eq!(
        byte_vec.as_ptr() as usize % std::mem::align_of::<f32>(),
        0,
        "Byte buffer is not aligned for f32"
    );

    // 安全地转换为f32切片
    let f32_slice = unsafe {
        std::slice::from_raw_parts(
            byte_vec.as_ptr() as *const f32,
            byte_vec.len() / size_of::<f32>()
        )
    };

    // 转换到u8并确保数值范围正确
    let image_data = f32_slice
        .iter()
        .map(|&f| (f.clamp(0.0, 1.0) * 255.0) as u8)
        .collect::<Vec<u8>>();

    // 创建并保存图像
    let image = GrayImage::from_raw(28, 28, image_data)
        .expect("Failed to create image from raw data");
    image.save("d:/abc.png").unwrap();
}
